import numpy as np
from collections import namedtuple
from config import *

Target = namedtuple('Target', 'name freq_list threat')

class JammingEnv:
    """
    极简版：6 个目标的频率集固定，每次给出所有频点。
    实际场景可用 STFT / 侦察结果动态更新 freq_list。
    """
    def __init__(self):
        self.targets = [
            Target('N1', np.random.uniform(310, 360, 32), 6),
            Target('N2', np.random.uniform(290, 330, 32), 5),
            Target('N3', np.random.uniform(250, 290, 32), 4),
            Target('N4', np.random.uniform(230, 270, 32), 3),
            Target('N5', np.random.uniform(220, 250, 32), 2),
            Target('N6', np.random.uniform(200, 230, 32), 1),
        ]
        self.all_freqs = np.concatenate([t.freq_list for t in self.targets])
        self.all_freqs.sort()

        # 按论文方法把频点划分子频段（贪心切分）
        self.sub_bands = self._split_sub_bands()
        self.n_sub = len(self.sub_bands)

        self.reset()

    # ---------- 环境接口 ----------
    def reset(self):
        self.jammed = [False] * len(self.targets)
        self.total_band = 0
        return self._state()

    def step(self, sub_idx, band_idx):
        """
        输入：
            sub_idx   子频段索引
            band_idx  带宽索引（0~10）
        返回：
            state, reward, done
        """
        bw = BAND_MIN + band_idx * BAND_STEP
        f_low, f_high = self.sub_bands[sub_idx]
        reward = 0

        # 统计该子频段内哪些目标被干扰
        for i, tgt in enumerate(self.targets):
            if self.jammed[i]:
                continue
            mask = (tgt.freq_list >= f_low) & (tgt.freq_list <= f_high)
            hit_ratio = mask.sum() / len(tgt.freq_list)
            # 论文：干扰 ≥1/3 频点即可
            if hit_ratio >= 1/3:
                self.jammed[i] = True
                reward += tgt.threat * 10   # 对应论文 r_base

        self.total_band += bw
        done = all(self.jammed)
        # 额外奖励：带宽利用率
        if reward > 0:
            reward += 10 if bw < 1.5 else 0
        else:
            reward = -2
        return self._state(), reward, done

    # ---------- 内部工具 ----------
    def _split_sub_bands(self):
        bands = []
        start = self.all_freqs[0]
        for f in self.all_freqs[1:]:
            if f - start > BAND_MAX:
                bands.append((start, start + BAND_MAX))
                start = f
        bands.append((start, self.all_freqs[-1]))
        return bands

    def _state(self):
        """
        返回两个层级的观测：
          s1: [每个子频段包含的频点个数]  -> length = n_sub
          s2: [sub_idx 对应的子频段内各目标频点占比] -> length = len(targets)
        """
        s1 = np.zeros(self.n_sub)
        for i, (low, high) in enumerate(self.sub_bands):
            s1[i] = ((self.all_freqs >= low) & (self.all_freqs <= high)).sum()

        # 占位，实际训练时再实时构造 s2
        return {"s1": s1.astype(np.float32),
                "s2_base": np.array([len(t.freq_list) for t in self.targets], np.float32)}

    # ---------- 可视化 ----------
    def summary(self):
        print("Targets jammed:", sum(self.jammed), "/", len(self.targets))
        print("Total bandwidth:", self.total_band, "MHz")